package com.fanshuai.tree.unionfind;

/**
 * 并查集
 * 并查集为树形的数据结构，表示多个不相交的集合
 * N个元素的集合，元素编号为0 ~ N-1
 * 使用指向父节点的树表示并查集，每个节点指向根节点
 *
 * 并查集可以解决多个集合的连通性问题，以及不相交集合个数等
 */
public class UnionFind implements IUnionFind {
    /**
     * 指向父节点的树  a[i]指向父节点的索引  根节点的父节点指向自身
     */
    private int[] a;
    /**
     * 集合元素个数  components[i]表示索引i元素所在的集合元素个数
     */
    private int[] components;
    /**
     * 元素个数
     */
    private int size;
    /**
     * 集合个数
     */
    private int componentSize;

    /**
     * 初始化
     * @param size  元素个数
     */
    public UnionFind(int size) {
        a = new int[size];
        this.size = size;
        this.componentSize = size;

        //初始化时，全部节点都为根节点
        for (int i = 0; i < size; i++) {
            a[i] = i;
            components[i] = 1;
        }
    }

    /**
     * 索引p和q指向的元素是否连通
     * @param p
     * @param q
     * @return
     */
    @Override
    public boolean isConnected(int p, int q) {
        int pRoot = find(p);
        int qRoot = find(q);

        return pRoot == qRoot;
    }

    /**
     * 连通2个元素所在的集合
     * 查找2个元素的根节点，比较集合大小，将较小的集合并入较大的集合
     * @param p
     * @param q
     */
    @Override
    public void unionElements(int p, int q) {
        int pRoot = find(p);
        int qRoot = find(q);

        if (pRoot == qRoot) {
            return;
        }

        /**
         * 将数量较小的集合加入较大的集合
         */
        if (components[pRoot] < components[qRoot]) {
            a[pRoot] = qRoot;
            components[qRoot] += components[pRoot];
        } else {
            a[qRoot] = pRoot;
            components[pRoot] += components[qRoot];
        }

        //元素数量-1
        componentSize--;
    }

    @Override
    public int size() {
        return size;
    }

    /**
     * 查询索引为p的元素所在的根节点
     * @param p
     * @return
     */
    @Override
    public int find(int p) {
        int root = p;

        //查询根节点
        while (a[root] != root) {
            root = a[root];
        }

        //路径压缩。将子节点的父亲全部指向根节点
        while (p != root) {
            int parent = a[p];
            a[p] = root;
            p = parent;
        }

        return root;
    }

    @Override
    public int componentSize(int p) {
        return components[find(p)];
    }

    @Override
    public int components() {
        return componentSize;
    }
}
